/*
 * pi3's Linux kernel Runtime Guard
 *
 * Component:
 *  - Red-black tree for keeping track of usermode processes
 *
 * Notes:
 *  - Makes sense with own kmem_cache_* allocation
 *
 * Timeline:
 *  - Created: 07.IX.2017
 *
 * Author:
 *  - Adam 'pi3' Zabrocki (http://pi3.com.pl)
 *
 */

#include "../../../../p_lkrg_main.h"


struct kmem_cache *p_ed_pids_cache = NULL;
struct p_tasks_root p_rb_hash[RB_HASH_SIZE] __attribute__ ((aligned(L1_CACHE_BYTES)));

struct p_ed_process *p_rb_find_ed_pid(struct rb_root *p_root, pid_t p_arg) {

   struct rb_node *p_node = p_root->rb_node;
   struct p_ed_process *p_struct = NULL;

   while(p_node) {
      p_struct = rb_entry(p_node, struct p_ed_process, p_rb);

      if (p_arg < p_struct->p_ed_task.p_pid) {
         p_node = p_node->rb_left;
      } else if (p_arg > p_struct->p_ed_task.p_pid) {
         p_node = p_node->rb_right;
      } else {
         return p_struct;
      }
   }

   return NULL;
}


struct p_ed_process *p_rb_add_ed_pid(struct rb_root *p_root, pid_t p_arg, struct p_ed_process *p_source) {

   struct rb_node **p_node = &p_root->rb_node;
   struct rb_node *p_parent = NULL;
   struct p_ed_process *p_struct;

   while(*p_node) {
      p_parent = *p_node;
      p_struct = rb_entry(p_parent, struct p_ed_process, p_rb);

      if (p_arg < p_struct->p_ed_task.p_pid) {
         p_node = &(*p_node)->rb_left;
      } else if (p_arg > p_struct->p_ed_task.p_pid) {
         p_node = &(*p_node)->rb_right;
      } else {
         return p_struct;
      }

   }
   rb_link_node(&p_source->p_rb, p_parent, p_node);   // Insert this new node as a red leaf
   rb_insert_color(&p_source->p_rb, p_root);          // Rebalance the tree, finish inserting

   return NULL;
}


void p_rb_del_ed_pid(struct rb_root *p_root, struct p_ed_process *p_source) {

   rb_erase(&p_source->p_rb, p_root);          // Erase the node
   p_reset_ed_flags(p_source);
   p_free_ed_pids(p_source);                   // Free the memory
}

static void p_ed_pids_cache_init(void *p_arg) {

   struct p_ed_process *p_struct = p_arg;

   memset(p_struct, 0, sizeof(struct p_ed_process));
}

int p_init_rb_ed_pids(void) {

   unsigned int i;

   for (i=0; i < RB_HASH_SIZE; i++) {
      p_rb_hash[i].p_tree.tree = RB_ROOT;
      rwlock_init(&p_rb_hash[i].p_lock.lock);
   }

   if ( (p_ed_pids_cache = kmem_cache_create("p_ed_pids", sizeof(struct p_ed_process),
                                           0, SLAB_HWCACHE_ALIGN, p_ed_pids_cache_init)) == NULL) {
      p_print_log(P_LKRG_ERR, "kmem_cache_create() for ED PIDs error! :(\n");
      return -ENOMEM;
   }

   return P_LKRG_SUCCESS;
}

void p_delete_rb_ed_pids(void) {

   struct rb_node *p_node;
   struct p_ed_process *p_tmp;
   unsigned int i;

   if (p_ed_pids_cache) {
      for (i=0; i<RB_HASH_SIZE; i++) {
         p_tasks_write_lock_noirq(&p_rb_hash[i].p_lock.lock);
         while ((p_node = rb_first(&p_rb_hash[i].p_tree.tree))) {
            p_tmp = rb_entry(p_node, struct p_ed_process, p_rb);
            p_print_log(P_LKRG_INFO, "Deleting ED PID => %d\n",p_tmp->p_ed_task.p_pid);
            p_rb_del_ed_pid(&p_rb_hash[i].p_tree.tree, p_tmp);

         }
         p_tasks_write_unlock_noirq(&p_rb_hash[i].p_lock.lock);
      }
      mb();
      kmem_cache_destroy(p_ed_pids_cache);
      p_ed_pids_cache = NULL;
   }

}
